package com.star.spring.ratelimit.limiter;

import java.util.Set;
import java.util.concurrent.TimeUnit;

import lombok.extern.slf4j.Slf4j;

import org.springframework.data.redis.core.RedisTemplate;
import org.springframework.stereotype.Component;

/**
 * Redis滑动窗口限流器
 *
 * @author star
 * @since 2025-01-24
 */
@Slf4j
@Component
public class RedisSlidingWindowLimiter {

    private final RedisTemplate<String, Object> redisTemplate;

    public RedisSlidingWindowLimiter(RedisTemplate<String, Object> redisTemplate) {
        this.redisTemplate = redisTemplate;
    }

    /**
     * 尝试获取请求许可
     *
     * @param key               限流key
     * @param maxRequests       最大请求数
     * @param windowSizeSeconds 窗口大小（秒）
     * @param precision         精度（秒）
     * @return 是否获取成功
     */
    public boolean tryAcquire(String key, int maxRequests, int windowSizeSeconds, int precision) {
        try {
            long currentTime = System.currentTimeMillis();
            long windowStart = currentTime - windowSizeSeconds * 1000L;

            // 使用ZSet存储请求时间戳
            String zsetKey = key + ":sliding";

            // 移除过期的记录
            redisTemplate.opsForZSet().removeRangeByScore(zsetKey, 0, windowStart);

            // 获取当前窗口内的请求数
            Long currentCount = redisTemplate.opsForZSet().count(zsetKey, windowStart, currentTime);

            if (currentCount != null && currentCount >= maxRequests) {
                log.debug("Redis sliding window limiter - key: {}, currentCount: {}, maxRequests: {}, acquired: false",
                        key, currentCount, maxRequests);
                return false;
            }

            // 添加当前请求
            redisTemplate.opsForZSet().add(zsetKey, currentTime, currentTime);

            // 设置过期时间
            redisTemplate.expire(zsetKey, windowSizeSeconds + precision, TimeUnit.SECONDS);

            log.debug("Redis sliding window limiter - key: {}, currentCount: {}, maxRequests: {}, acquired: true",
                    key, (currentCount != null ? currentCount : 0) + 1, maxRequests);
            return true;

        } catch (Exception e) {
            log.error("Redis sliding window limiter error - key: {}", key, e);
            // Redis异常时允许通过
            return true;
        }
    }

    /**
     * 尝试获取请求许可（使用默认参数）
     *
     * @param key 限流key
     * @return 是否获取成功
     */
    public boolean tryAcquire(String key) {
        return tryAcquire(key, 1000, 60, 1);
    }

    /**
     * 获取当前窗口的请求数
     *
     * @param key               限流key
     * @param windowSizeSeconds 窗口大小（秒）
     * @return 当前请求数
     */
    public int getCurrentCount(String key, int windowSizeSeconds) {
        try {
            long currentTime = System.currentTimeMillis();
            long windowStart = currentTime - windowSizeSeconds * 1000L;

            String zsetKey = key + ":sliding";
            Long count = redisTemplate.opsForZSet().count(zsetKey, windowStart, currentTime);
            return count != null ? count.intValue() : 0;
        } catch (Exception e) {
            log.error("Get current count error - key: {}", key, e);
            return 0;
        }
    }

    /**
     * 获取窗口内的所有请求时间戳
     *
     * @param key               限流key
     * @param windowSizeSeconds 窗口大小（秒）
     * @return 请求时间戳集合
     */
    public Set<Object> getWindowRequests(String key, int windowSizeSeconds) {
        try {
            long currentTime = System.currentTimeMillis();
            long windowStart = currentTime - windowSizeSeconds * 1000L;

            String zsetKey = key + ":sliding";
            return redisTemplate.opsForZSet().rangeByScore(zsetKey, windowStart, currentTime);
        } catch (Exception e) {
            log.error("Get window requests error - key: {}", key, e);
            return null;
        }
    }
}
